"""
create_train_yaml.py
create by cuiyu 2025-08-05
说明：
 * 说明内容

用法:
  create_train_yaml.py

更新:
2025-08-05:
"""

import sys
# import logging
from pathlib import Path
from ruamel.yaml import YAML

# _APP_FILE = Path(sys.argv[0])
# _APP_ROOT = _APP_FILE.parent

def CreateTrainYaml(data_root: Path, class_ids: dict):
    """
    生成 train.yaml 文件

    Args:
        :data_root: 数据集根目录， == img_root.parent
        :class_ids: 类别字典，{id, name} 
    """
    # 设置 train_file
    train_file = data_root / "train.yaml"

    # 准备 yaml 数据
    yaml_data: dict = {}
    yaml_data["path"] = str(data_root)
    yaml_data["train"] = "train.txt"
    yaml_data["val"] = "val.txt"
    yaml_data["nc"] = len(class_ids)
    yaml_data["names"] = class_ids

    # 创建 YAML 对象
    yaml = YAML(typ="rt")
    yaml.preserve_quotes = True  # 保留字符串引号（可选）
    yaml.indent(mapping=2, sequence=4, offset=2)  # 设置缩进格式
    yaml.allow_unicode = True  # 关键：禁用 Unicode 转义

    # 写入 train.yaml 内容
    with open(train_file, "w", encoding="utf-8") as fwrite:
        yaml.dump(yaml_data, fwrite)

def _test():
    from bhclasses_define import BHClasses

    img_root: Path = Path("/aaby/netdisk/mydata/ai-data/dataset/priv/井盖病害/训练数据/img01")
    CreateTrainYaml(img_root.parent, BHClasses.dict_id_name)
    pass

if __name__ == "__main__":
    # logging.basicConfig(level=logging.DEBUG) 
    _test()